package org.spider.download.httpclient;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.apache.http.HttpResponse;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.util.EntityUtils;
import org.spider.download.Downloader;
import org.spider.model.DownloadRequest;
import org.spider.model.DownloadResponse;
import org.spider.util.HttpUtils;
import org.spider.utils.CharsetUtils;

import java.io.IOException;
import java.nio.charset.Charset;

/**
 * The http downloader based on HttpClient.
 *
 * @author code4crafter@gmail.com <br>
 * @since 0.1.0
 */
@Slf4j
@Data
public class HttpClientDownloader implements Downloader {

    private HttpClientGenerator httpClientGenerator = new HttpClientGenerator();

    private HttpUriRequestConverter httpUriRequestConverter = new HttpUriRequestConverter();

    private CloseableHttpClient getHttpClient(DownloadRequest downloadRequest) {
        return httpClientGenerator.getClient(downloadRequest);
    }

    @Override
    public DownloadResponse download(DownloadRequest downloadRequest) {
        CloseableHttpResponse httpResponse = null;
        CloseableHttpClient httpClient = getHttpClient(downloadRequest);
        HttpClientRequestContext requestContext = httpUriRequestConverter.convert(downloadRequest);
        DownloadResponse downloadResponse = new DownloadResponse();
        try {
            httpResponse = httpClient.execute(requestContext.getHttpUriRequest(), requestContext.getHttpClientContext());
            downloadResponse = handleResponse(downloadRequest, downloadRequest.getCharset(), httpResponse);
            return downloadResponse;
        } catch (IOException e) {
            log.error("download page {} error", downloadRequest.getUrl(), e);
            return downloadResponse;
        } finally {
            if (httpResponse != null) {
                //ensure the connection is released back to pool
                EntityUtils.consumeQuietly(httpResponse.getEntity());
            }
        }
    }

    protected DownloadResponse handleResponse(DownloadRequest downloadRequest, String charset, HttpResponse httpResponse) throws IOException {
        byte[] bytes = IOUtils.toByteArray(httpResponse.getEntity().getContent());
        String contentType = httpResponse.getEntity().getContentType() == null ? "" : httpResponse.getEntity().getContentType().getValue();
        DownloadResponse downloadResponse = new DownloadResponse();
        downloadResponse.setBytes(bytes);
        if (!downloadRequest.isBinaryContent()){
            if (charset == null) {
                charset = getHtmlCharset(contentType, bytes);
            }
            downloadResponse.setCharset(charset);
            downloadResponse.setRawText(new String(bytes, charset));
        }
        downloadResponse.setDownloadRequest(downloadRequest);
        downloadResponse.setStatusCode(httpResponse.getStatusLine().getStatusCode());
        downloadResponse.setHeaderMap(HttpUtils.convertHeaders(httpResponse.getAllHeaders()));
        return downloadResponse;
    }

    private String getHtmlCharset(String contentType, byte[] contentBytes) throws IOException {
        String charset = CharsetUtils.detectCharset(contentType, contentBytes);
        if (charset == null) {
            charset = Charset.defaultCharset().name();
            log.warn("Charset autodetect failed, use {} as charset. Please specify charset in Site.setCharset()", Charset.defaultCharset());
        }
        return charset;
    }
}
